#!/usr/bin/env python
# -*- coding:utf-8 -*-

import pefile,collections
import traceback
class PeInfo(object):
    def __init__(self,fpath):
        self.fpath = fpath
    def load_pe(self):
        try:
            self.pe = pefile.PE(data=self.fpath)
            self.imageBase = self.pe.OPTIONAL_HEADER.ImageBase
            self.addrEntry = self.imageBase + self.pe.OPTIONAL_HEADER.AddressOfEntryPoint
            if self.pe.FILE_HEADER.Machine == 0x014c:#i386即32位标志
                self.arch = 32
            elif self.pe.FILE_HEADER.Machine == 0x8664:#amd64即64位标志
                self.arch = 64
        except Exception as e:
            print("PE文件加载失败!")
            traceback.print_exc()
            return False
        return True
    def get_imageBase(self):
        return self.imageBase
    def get_pefile(self):
        return self.pe
    def get_addrEntry(self):
        return self.addrEntry
    def get_arch(self):
        return self.arch
    def get_SectionAlignment(self):
        sa = self.pe.OPTIONAL_HEADER.SectionAlignment
        return sa
    def get_sections(self):
        return self.pe.sections
    def get_data(self,rva,length):
        data = self.pe.get_data(rva,length)
        return data
    def get_iat(self):
        iat = self.pe.DIRECTORY_ENTRY_IMPORT
        return iat
